fix: EnforceDistribution optimizer preserves fetch (LIMIT) from distribution-changing operators by zhuqi-lucas · Pull Request #21170 · apache/datafusion

zhuqi-lucas · 2026-03-26T08:57:56Z

Which issue does this PR close?

Closes Bug: EnforceDistribution optimizer loses fetch (LIMIT) from CoalescePartitionsExec and SortPreservingMergeExec #21169.

Rationale for this change

When LimitPushdown merges a GlobalLimitExec into a CoalescePartitionsExec or SortPreservingMergeExec as a fetch value, the EnforceDistribution optimizer rule strips and re-inserts these distribution-changing operators without preserving the fetch. This silently drops the LIMIT for queries over multi-partition sources, potentially returning duplicate/extra rows.

What changes are included in this PR?

remove_dist_changing_operators now captures any fetch value and the LexOrdering from a stripped SortPreservingMergeExec (if present) before stripping operators. Only the ordering is captured — not the full plan node — so that a fresh SPM can be reconstructed later with the current (possibly rewritten) child plan, avoiding stale references.
add_merge_on_top accepts an optional fetch and applies it to the newly created SortPreservingMergeExec or CoalescePartitionsExec.
Fallback re-introduction: if the fetch was not consumed by add_merge_on_top (e.g., when the parent had UnspecifiedDistribution or the child already had a single partition), the limit is re-introduced as a wrapping operator so it is never silently lost. When the original operator was an SPM, a fresh SortPreservingMergeExec is reconstructed from the captured ordering to preserve sort semantics.
SPM capture decoupled from fetch: the ordering capture uses spm_ordering.is_none() instead of gating on fetch.is_none(), so the SPM ordering is recorded even when an outer operator (e.g. CoalescePartitionsExec) already set fetch.

Are these changes tested?

Yes, four tests are added:

coalesce_partitions_fetch_preserved_by_enforce_distribution — unsorted multi-partition source with CoalescePartitionsExec(fetch=1)
coalesce_partitions_fetch_preserved_sorted — sorted multi-partition source with CoalescePartitionsExec(fetch=5)
spm_fetch_preserved_by_enforce_distribution — sorted multi-partition source with SortPreservingMergeExec(fetch=3)
nested_coalesce_over_spm_preserves_spm_ordering — nested CoalescePartitionsExec(fetch=5) over SortPreservingMergeExec(fetch=3), verifying that SPM ordering is preserved even when an outer operator already set fetch

Are there any user-facing changes?

No API changes. Queries with LIMIT over multi-partition sources will now correctly preserve the limit through the EnforceDistribution optimizer pass.

Copilot

Pull request overview

Fixes a correctness bug in the physical optimizer where EnforceDistribution could drop a fetch (LIMIT) that had been pushed down into distribution-changing operators (CoalescePartitionsExec / SortPreservingMergeExec), leading to extra/duplicate rows for multi-partition inputs.

Changes:

Track and propagate fetch while stripping distribution-changing operators so it can be re-applied later.
Extend add_merge_on_top to optionally apply fetch when re-adding SortPreservingMergeExec.
Add regression tests ensuring fetch survives EnforceDistribution across coalesce and SPM cases.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
`datafusion/physical-optimizer/src/enforce_distribution.rs`	Captures `fetch` while stripping dist-changing ops; attempts to reapply it during merge insertion or via a fallback wrapper.
`datafusion/core/tests/physical_optimizer/enforce_distribution.rs`	Adds regression tests to ensure `fetch` is preserved through `EnforceDistribution`.

Comments suppressed due to low confidence (1)

datafusion/physical-optimizer/src/enforce_distribution.rs:992

add_merge_on_top only applies the stripped fetch to a newly created SortPreservingMergeExec, but not to the CoalescePartitionsExec branch. If the child has no ordering, this leaves fetch unconsumed and triggers the fallback wrapper later, potentially producing redundant CoalescePartitionsExec nodes and extra overhead. Consider applying fetch.take() to CoalescePartitionsExec::new(...).with_fetch(...) as well when adding the merge/coalesce on top.

        let new_plan = if let Some(req) = input.plan.output_ordering() {
            Arc::new(
                SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan))
                    .with_fetch(fetch.take()),
            ) as _
        } else {
            // If there is no input order, we can simply coalesce partitions:
            Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) as _
        };

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

rkrishn7

Thanks @zhuqi-lucas! Left a couple comments

rkrishn7 · 2026-03-26T16:51:59Z

+    if let Some(fetch_val) = fetch.take() {
+        let limit_plan: Arc<dyn ExecutionPlan> = if let Some(spm) = spm {
+            // Re-insert the original SortPreservingMergeExec with fetch.
+            spm.with_fetch(Some(fetch_val)).unwrap()


Hmm, is spm actually stale here? It was captured prior to processing of the plan's children above. I think we might want to use the SortPreservingMergeExec constructor directly here with the new plan.

Good catch! Yes, spm was indeed stale — it was captured in remove_dist_changing_operators before the child subtree went through ensure_distribution, so it referenced the old (pre-rewrite) child plan.

Fixed in 8d2ee25: instead of capturing the entire Arc<dyn ExecutionPlan>, we now only capture the LexOrdering from the SPM. In the fallback path, we reconstruct a fresh SortPreservingMergeExec::new(ordering, dist_context.plan) using the current (rewritten) child. This also eliminates the unwrap().

Additionally, the spm capture was gated on fetch.is_none(), which meant it would be skipped when an outer operator (e.g. CoalescePartitionsExec) already set fetch. I decoupled it to use spm_ordering.is_none() instead. Added a regression test (nested_coalesce_over_spm_preserves_spm_ordering) that confirms the old code produced CoalescePartitionsExec (losing sort semantics) while the fix correctly produces SortPreservingMergeExec.

rkrishn7 · 2026-03-26T16:54:32Z

+        } else {
+            // The fetch came from a CoalescePartitionsExec. Re-introduce
+            // it as a CoalescePartitionsExec(fetch=N) wrapping the output.
+            Arc::new(


Do we want to check if the new plan outputs more than a single partition? And if so, use GlobalLimitExec instead of CoalescePartitionsExec?

Not sure this one matters too much though.

Good point. In the current flow, when fetch is unconsumed it means add_merge_on_top didn't fire (either UnspecifiedDistribution or single partition). For single partition the CoalescePartitionsExec is essentially a no-op wrapper with fetch. For multi-partition cases, the fetch was already consumed by add_merge_on_top. So in practice I think the current behavior is correct, but I'll keep an eye on edge cases. Thanks for raising it!

…ibution-changing operators When `LimitPushdown` merges a `GlobalLimitExec` into a `CoalescePartitionsExec` or `SortPreservingMergeExec` as a `fetch` value, `EnforceDistribution` would strip and re-insert these operators without preserving the fetch. This silently drops the LIMIT for queries over multi-partition sources. The fix captures the `fetch` in `remove_dist_changing_operators` and threads it through `add_merge_on_top`, or re-introduces it as a fallback when the merge operator is not re-inserted. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…re from fetch - Capture only the LexOrdering from the stripped SortPreservingMergeExec instead of the entire plan node, then reconstruct a fresh SPM with the current (possibly rewritten) child plan in the fallback path. This avoids reusing a stale SPM that references an outdated subtree. - Decouple spm_ordering capture from fetch — use `spm_ordering.is_none()` instead of gating on `fetch.is_none()`, so the SPM ordering is recorded even when an outer operator already set fetch. - Remove unwrap() on with_fetch by constructing SPM directly. - Fix test comment: fallback is CoalescePartitionsExec/SPM, not GlobalLimitExec. - Add regression test: nested CoalescePartitionsExec(fetch=5) over SortPreservingMergeExec(fetch=3) now correctly preserves SPM ordering.

github-actions · 2026-06-27T02:17:37Z

Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days.

Copilot AI review requested due to automatic review settings March 26, 2026 08:57

github-actions Bot added optimizer Optimizer rules core Core DataFusion crate labels Mar 26, 2026

Copilot started reviewing on behalf of zhuqi-lucas March 26, 2026 08:58 View session

Copilot AI reviewed Mar 26, 2026

View reviewed changes

Comment thread datafusion/physical-optimizer/src/enforce_distribution.rs Outdated

Comment thread datafusion/core/tests/physical_optimizer/enforce_distribution.rs Outdated

Comment thread datafusion/physical-optimizer/src/enforce_distribution.rs Outdated

zhuqi-lucas force-pushed the fix/enforce-dist-preserve-fetch branch 2 times, most recently from 836549e to c8cc0b3 Compare March 26, 2026 09:21

zhuqi-lucas mentioned this pull request Mar 26, 2026

Bug: EnforceDistribution is not idempotent despite documentation claim #21174

Open

rkrishn7 reviewed Mar 26, 2026

View reviewed changes

zhuqi-lucas force-pushed the fix/enforce-dist-preserve-fetch branch from 8d2ee25 to 5b0a42e Compare March 27, 2026 06:41

zhuqi-lucas and others added 3 commits March 27, 2026 14:45

fmt

493d1bf

zhuqi-lucas force-pushed the fix/enforce-dist-preserve-fetch branch from 5b0a42e to d861c98 Compare March 27, 2026 06:47

Merge branch 'main' into fix/enforce-dist-preserve-fetch

b7dc8ef

github-actions Bot added the Stale PR has not had any activity for some time label Jun 27, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: EnforceDistribution optimizer preserves fetch (LIMIT) from distribution-changing operators#21170

fix: EnforceDistribution optimizer preserves fetch (LIMIT) from distribution-changing operators#21170
zhuqi-lucas wants to merge 4 commits into
apache:mainfrom
zhuqi-lucas:fix/enforce-dist-preserve-fetch

zhuqi-lucas commented Mar 26, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rkrishn7 left a comment

Uh oh!

rkrishn7 Mar 26, 2026

Uh oh!

zhuqi-lucas Mar 27, 2026

Uh oh!

rkrishn7 Mar 26, 2026

Uh oh!

zhuqi-lucas Mar 27, 2026

Uh oh!

github-actions Bot commented Jun 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

zhuqi-lucas commented Mar 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

rkrishn7 left a comment

Choose a reason for hiding this comment

Uh oh!

rkrishn7 Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

zhuqi-lucas Mar 27, 2026

Choose a reason for hiding this comment

Uh oh!

rkrishn7 Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

zhuqi-lucas Mar 27, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jun 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zhuqi-lucas commented Mar 26, 2026 •

edited

Loading